import pandas as pd
from math import log, exp, pow, sqrt, pi

def class_num(data, index):
    """
    输出第index个属性可能的取值数
    """
    class_number = {}
    for row in data:
        if row[index] not in class_number.keys():
            class_number[row[index]] = 0
        class_number[row[index]] += 1

    return len(class_number)


def continuous_attribute(data, labels, label, index):
    """
    连续属性的均值和方差
    """
    ave = 0.0  # 均值
    var = 0.0  # 方差
    count = 0

    # 计算label类对应的均值
    for row in range(len(data)):
        if labels[row] == label:
            count += 1
            ave += data[row, index]
    ave = ave / count

    # 计算label类对应的方差
    for row in range(len(data)):
        if labels[row] == label:
            var += (data[row, index] - ave) * (data[row, index] - ave)
    var = var / count

    return ave, var


if __name__ == '__main__':
    # 读取数据
    df = pd.read_csv('watermelon_4_3.csv')
    data = df.values[:, 1:-1]
    labels = df.values[:, -1].tolist()

    test = df.values[0, 1:-1]  # “测1”样本

    # 估计类先验概率（使用对数似然将连乘变为连加）
    prob_good = log((8 + 1) / (17 + 2))  # 公式(7.19)
    prob_bad = log((9 + 1) / (17 + 2))  # 公式(7.19)

    for i in range(len(data[0])):
        # 为每个属性估计条件概率
        if type(test[i]).__name__ == 'float':
            # 该属性是连续值
            ave0, var0 = continuous_attribute(data, labels, 0, i)  # 坏瓜在该属性上的均值和方差
            ave1, var1 = continuous_attribute(data, labels, 1, i)  # 好瓜在该属性上的均值和方差
            prob0 = exp(-pow(test[i] - ave0, 2) / (2 * var0)) / sqrt(2 * pi * var0)  # 公式(7.18)
            prob1 = exp(-pow(test[i] - ave1, 2) / (2 * var1)) / sqrt(2 * pi * var1)  # 公式(7.18)
            prob_good += log(prob1)  # 公式(7.15)
            prob_bad += log(prob0)  # 公式(7.15)
        
        else:
            # 该属性是离散值
            count_good = 0
            count_bad = 0
            for row in range(len(data)):
                if data[row, i] == test[i]:
                    if labels[row] == 1:
                        count_good += 1
                    if labels[row] == 0:
                        count_bad += 1
            
            prob_good += log((count_good + 1) / (8 + class_num(data, i)))  # 公式(7.20)
            prob_bad += log((count_bad + 1) / (9 + class_num(data, i)))  # 公式(7.20)

    print('在“测1”样本上，好瓜的对数似然：', prob_good)
    print('在“测1”样本上，坏瓜的对数似然：', prob_bad)
    if prob_good >= prob_bad:
        print('“测1”的预测结果为：好瓜')
    else:
        print('“测1”的预测结果为：坏瓜')
